"""
书本检测模块，使用ultralytics YOLO模型检测图像中的书本。
"""
import os
import cv2
import numpy as np
from typing import List, Tuple, Dict, Optional, Union
from dataclasses import dataclass
from ultralytics import YOLO


@dataclass
class DetectionResult:
    """检测结果数据类"""
    box: Tuple[int, int, int, int]  # 左上x, 左上y, 宽度, 高度
    confidence: float
    class_id: int
    class_name: str


class BookDetector:
    """
    书本检测器类，使用ultralytics YOLO模型检测图像中的书本。
    """

    def __init__(
        self,
        model_path: Optional[str] = None,
        confidence_threshold: float = 0.5,
        device: Optional[str] = None
    ):
        """
        初始化书本检测器

        参数:
            model_path: YOLO模型路径，如果为None则使用默认路径
            confidence_threshold: 置信度阈值
            device: 运行设备，可以是'cuda'或'cpu'，如果为None则自动选择
        """
        self.confidence_threshold = confidence_threshold

        # 设置运行设备
        if device is None:
            import torch
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device

        # 加载模型
        if model_path is None:
            # 如果未指定模型路径，则查找默认路径
            default_paths = [
                os.path.join(os.path.dirname(__file__),
                             '../models/yolov11_book.pt'),
                os.path.join(os.path.dirname(__file__),
                             '../../models/yolov11_book.pt'),
                'yolov11_book.pt'
            ]

            for path in default_paths:
                if os.path.exists(path):
                    model_path = path
                    break

        try:
            # 尝试加载指定的模型
            if model_path is not None and os.path.exists(model_path):
                self.model = YOLO(model_path)
                print(f"成功加载YOLO模型 {model_path}，运行设备: {self.device}")
            else:
                # 如果没有找到指定模型，使用预训练模型
                print("未找到指定YOLO模型，加载预训练模型...")
                self.model = YOLO("yolov8n.pt")  # 加载小型预训练模型
                print(f"已加载预训练YOLO模型，运行设备: {self.device}")

            # 配置模型运行设备
            self.model = self.model.to(self.device)

            # 获取类别列表
            self.class_names = self.model.names if hasattr(
                self.model, 'names') else {}
            print(f"检测类别: {self.class_names}")

        except Exception as e:
            # 如果加载失败，使用模拟检测
            print(f"加载YOLO模型失败: {str(e)}，将使用模拟检测")
            self.model = None

    def detect_books(self, image: np.ndarray) -> List[DetectionResult]:
        """
        检测图像中的书本

        参数:
            image: 输入图像

        返回:
            List[DetectionResult]: 检测结果列表
        """
        if image is None or image.size == 0:
            return []

        # 如果模型为None，使用模拟检测
        if self.model is None:
            return self._simulate_detection(image)

        # 使用YOLO进行检测
        results = self.model(
            image, conf=self.confidence_threshold, verbose=False)

        # 处理检测结果
        detections = []

        # 提取检测结果 - 使用ultralytics API
        for result in results:
            boxes = result.boxes
            for i, box in enumerate(boxes):
                try:
                    # 获取边界框
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

                    # 获取置信度
                    conf = float(box.conf[0].cpu().numpy())

                    # 获取类别ID和名称
                    cls = int(box.cls[0].cpu().numpy())
                    cls_name = self.class_names.get(cls, f"class_{cls}")

                    # 只保留书本相关类别
                    if self._is_book_class(cls_name, cls):
                        detections.append(DetectionResult(
                            box=(x1, y1, x2-x1, y2-y1),
                            confidence=conf,
                            class_id=cls,
                            class_name=cls_name
                        ))
                except Exception as e:
                    print(f"处理检测框时出错: {str(e)}")
                    continue

        # 如果没有检测到书本，尝试使用模拟检测
        if not detections:
            print("YOLO未检测到书本，尝试使用模拟检测...")
            detections = self._simulate_detection(image)

        return detections

    def _is_book_class(self, class_name: str, class_id: int) -> bool:
        """判断是否为书本相关类别"""
        book_classes = {'book', 'textbook', '书',
                        '课本', 'books', 'notebook', 'paper'}
        book_class_ids = {0, 73}  # COCO数据集中的book类别ID为73

        return (class_name.lower() in book_classes or
                '书' in class_name or
                class_id in book_class_ids)

    def _simulate_detection(self, image: np.ndarray) -> List[DetectionResult]:
        """
        当无法使用YOLO模型时，使用传统计算机视觉方法模拟检测
        """
        h, w = image.shape[:2]
        total_area = h * w

        # 转为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 高斯模糊减少噪声
        blurred = cv2.GaussianBlur(gray, (7, 7), 0)

        # 边缘检测
        edges = cv2.Canny(blurred, 50, 150)

        # 膨胀边缘
        kernel = np.ones((5, 5), np.uint8)
        dilated = cv2.dilate(edges, kernel, iterations=2)

        # 查找轮廓
        contours, _ = cv2.findContours(
            dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        detections = []
        min_area = 10000  # 最小书本面积
        max_area_ratio = 0.8  # 最大书本面积占比

        # 对每个轮廓进行处理
        for contour in contours:
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            area_ratio = area / total_area

            # 过滤太小或太大的区域
            if area < min_area or area_ratio > max_area_ratio:
                continue

            # 获取边界框
            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = w / h if h > 0 else 0

            # 过滤不符合书本形状的区域（书本通常是矩形，宽高比在一定范围内）
            if not (0.3 < aspect_ratio < 3.0):
                continue

            # 计算置信度（基于形状和大小）
            confidence = min(0.8, 0.5 + 0.5 * area_ratio)

            # 创建检测结果
            detection = DetectionResult(
                box=(x, y, w, h),
                confidence=confidence,
                class_id=0,
                class_name="book"
            )

            detections.append(detection)

        return detections

    def get_book_regions(self, image: np.ndarray) -> List[Tuple[np.ndarray, Tuple[int, int, int, int]]]:
        """
        获取书本区域的图像和对应的边界框

        参数:
            image: 输入图像

        返回:
            List[Tuple[np.ndarray, Tuple[int, int, int, int]]]: 书本区域图像和边界框列表
        """
        detections = self.detect_books(image)

        book_regions = []
        for detection in detections:
            x, y, w, h = detection.box
            book_image = image[y:y+h, x:x+w].copy()
            book_regions.append((book_image, detection.box))

        return book_regions

    def draw_detections(self, image: np.ndarray, detections: List[DetectionResult], color=(0, 255, 0), thickness=2) -> np.ndarray:
        """
        在图像上绘制检测结果

        参数:
            image: 输入图像
            detections: 检测结果列表
            color: 边界框颜色
            thickness: 边界框线条粗细

        返回:
            np.ndarray: 绘制了检测结果的图像
        """
        result_image = image.copy()

        for detection in detections:
            x, y, w, h = detection.box
            cv2.rectangle(result_image, (x, y), (x+w, y+h), color, thickness)

            # 显示类别和置信度
            label = f"{detection.class_name}: {detection.confidence:.2f}"
            cv2.putText(
                result_image,
                label,
                (x, y-10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                thickness
            )

        return result_image
